import os
import argparse
from _src.load_save_model import loadcaffemodel, saveonnxmodel
from _src.caffe2onnx import Caffe2Onnx

def make_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default="./net_inst.txt", type=str, required=True, help='config file for inference')
    return parser

if __name__ == '__main__':
    myparser = make_parser()
    args = myparser.parse_args()
    config_file = args.config
    file_config = open(args.config).read()
    lines_config = file_config.split("\n")
    for i in range(0,len(lines_config)):
        if "prototxt_file" in lines_config[i]:
            prototxt_file=lines_config[i].split("=")[-1]
        elif "caffemodel_file" in lines_config[i]:
            caffemodel_file = lines_config[i].split("=")[-1]
        elif "inst_file" in lines_config[i]:
            inst_file = lines_config[i].split("=")[-1]
            onnx_file_name = inst_file.split("/")[-1]
            inst_file += ".onnx"


    graph, params = loadcaffemodel(prototxt_file,caffemodel_file)
    print("onnx name : ", onnx_file_name)
    c2o = Caffe2Onnx(graph, params, onnx_file_name)
    onnxmodel = c2o.createOnnxModel()
    saveonnxmodel(onnxmodel, inst_file)
